import os
import re
import numpy as np
import matplotlib.pyplot as plt


def parse_log_file(file_path):
    pattern = re.compile(r'Average reward score:\s+([\d.]+)')  # 只匹配奖励值
    rewards = []
    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
        for line in f:
            match = pattern.search(line)
            if match:
                reward = float(match.group(1))  # 提取奖励值
                rewards.append(reward)
    return rewards


def process_folder(folder_path):
    log_files = [f for f in os.listdir(folder_path) if f.endswith('.log')]
    all_rewards = []
    for log_file in log_files:
        file_path = os.path.join(folder_path, log_file)
        rewards = parse_log_file(file_path)
        if rewards:
            all_rewards.append(rewards)
    if not all_rewards:
        print(f"No valid data found in folder: {folder_path}")
        return None

    # 找到最短的奖励序列长度
    min_length = min(len(rewards) for rewards in all_rewards)

    # 截取所有奖励序列到最短长度
    truncated_rewards = [rewards[:min_length] for rewards in all_rewards]

    # 计算统计量
    mean_values = []
    max_values = []
    min_values = []
    for i in range(min_length):
        step_rewards = [rewards[i] for rewards in truncated_rewards]
        mean_values.append(np.mean(step_rewards))
        max_values.append(np.max(step_rewards))
        min_values.append(np.min(step_rewards))

    print(f"Processed folder: {folder_path}")
    print(f"Total steps: {min_length}")
    print(f"Mean values: {mean_values[:5]}...")
    print(f"Max values: {max_values[:5]}...")
    print(f"Min values: {min_values[:5]}...")

    return {
        'folder_name': os.path.basename(folder_path),
        'steps': list(range(1, min_length + 1)),  # 步数从 1 开始
        'mean': mean_values,
        'max': max_values,
        'min': min_values
    }


def plot_multiple_folders(folder_paths):
    all_data = []
    for folder in folder_paths:
        data = process_folder(folder)
        if data is not None and len(data['steps']) > 0:  # 检查数据是否为空
            all_data.append(data)

    if not all_data:  # 如果没有数据，提示用户
        print("No valid data found in any folder.")
        return

    plt.figure(figsize=(12, 6))
    for data in all_data:
        steps = data['steps']
        mean = data['mean']
        max_vals = data['max']
        min_vals = data['min']
        # 绘制平均线，并设置标签
        line, = plt.plot(steps[::5], mean[::5], label=data['folder_name'], lw=1.5)
        color = line.get_color()
        print(color)
        # 填充最大值和最小值区域
        plt.fill_between(steps, min_vals, max_vals, color=color, alpha=0.2, label=None)

    plt.xlabel('Steps', fontsize=12)
    plt.ylabel('Average Reward Score', fontsize=12)
    plt.title('Average Reward with Min-Max Range Across Experiments', fontsize=14)
    plt.legend(loc='best', fontsize=10)  # 确保图例被正确添加
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    # 修改此处为你的文件夹路径列表
    folders_to_plot = [
        r'C:\Users\Admin\Desktop\RLHFstep3_traininglog\RLHFstep3_traininglog\第二次实验\ABC',
        r'C:\Users\Admin\Desktop\RLHFstep3_traininglog\RLHFstep3_traininglog\第二次实验\uniform',
        r'C:\Users\Admin\Desktop\RLHFstep3_traininglog\RLHFstep3_traininglog\第二次实验\myrlhf',
        r'C:\Users\Admin\Desktop\RLHFstep3_traininglog\RLHFstep3_traininglog\第二次实验\rlhf'
    ]
    plot_multiple_folders(folders_to_plot)